package Algorithm;

import Structure.Edge;
import Structure.Graph;
import Structure.GraphGenerator;
import Structure.Node;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * @Author: ylc
 * @Description:prim算法
 * @Date Created in 2022-06-07 16:44
 */
public class Prim {
    public static class EdgeComparator implements Comparator<Edge> {
        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }

    /**
     * 普利姆最小生成树算法算法
     * @param graph 无向图
     * @return
     */
    public static Set<Edge> primMST(Graph graph){
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(
                new EdgeComparator());
        HashSet<Node> set = new HashSet<>();
        Set<Edge> result = new HashSet<>();
        for (Node node :    // 可以解决森林问题
                graph.nodes.values()) {
            if (!set.contains(node)){
                set.add(node);
                for (Edge edge :
                        node.edges) {
                    priorityQueue.offer(edge);
                }
                while(!priorityQueue.isEmpty()){
                    Edge edge = priorityQueue.poll();
                    Node toNode = edge.to;
                    if(!set.contains(toNode)){
                        set.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge :
                                toNode.edges) {
                            priorityQueue.offer(nextEdge);
                        }
                    }
                }
            }
        }
        return result;
    }
    public static void main(String[] args) {
        Graph graph = create();
        Set<Edge> edges = primMST(graph);
        for (Edge edge :
                edges) {
            System.out.print(edge.weight + " ");
        }
    }
    public static Graph create(){
        int[][] arr = {
                {6,1,2},{6,2,1},
                {1,1,3},{1,3,1},
                {5,1,4},{5,4,1},
                {5,2,3},{5,3,2},
                {5,3,4},{5,4,3},
                {3,2,5},{3,5,2},
                {6,3,5},{6,5,3},
                {4,3,6},{4,6,3},
                {2,4,6},{2,6,4},
                {6,5,6},{6,6,5},
        };

        return GraphGenerator.generateGraph(arr);
    }

}
